#!/usr/bin/env python3
"""
MCP 数学运算服务器
提供基本的数学运算功能：加减乘除乘方等
使用stdio传输协议与客户端通信
"""

from mcp.server.fastmcp import FastMCP
import asyncio
import math
import numpy as np
import scipy.optimize as opt
import scipy.interpolate as interp
import scipy.integrate as integrate
from scipy.stats import jarque_bera, shapiro
from sklearn.cluster import KMeans, DBSCAN, AgglomerativeClustering
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.tree import DecisionTreeClassifier
from sklearn.metrics import accuracy_score, f1_score, confusion_matrix, silhouette_score
from sklearn.preprocessing import StandardScaler
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import plotly.graph_objects as go
import plotly.express as px
from statsmodels.tsa.arima.model import ARIMA
from statsmodels.tsa.stattools import adfuller
from statsmodels.graphics.tsaplots import plot_acf, plot_pacf
import warnings
import base64
import io
from typing import Union, List, Dict, Any, Tuple, Optional
warnings.filterwarnings('ignore')

# 创建MCP服务器实例
mcp = FastMCP("mcp-basic-optimization")

# ==================== 非线性规划 - 无约束优化模块 ====================

@mcp.tool()
def newton_optimization(objective_func: str, initial_point: List[float], 
                       tolerance: float = 1e-6, max_iterations: int = 100, minimize: bool = True) -> Union[Dict[str, Any], str]:
    """
    牛顿法无约束优化
    
    参数:
    - objective_func: 目标函数表达式
    - initial_point: 初始点
    - tolerance: 收敛容差
    - max_iterations: 最大迭代次数
    - minimize: True为求最小值，False为求最大值
    """
    try:
        def objective(x):
            x = np.asarray(x, dtype=float)
            x_val = float(x[0]) if len(x) > 0 else 0.0
            y_val = float(x[1]) if len(x) > 1 else 0.0
            if objective_func == "(x-2)**2 + (y-1)**2":
                result = (x_val - 2.0)**2 + (y_val - 1.0)**2
            else:
                obj_expr = objective_func.replace('x', str(x_val)).replace('y', str(y_val))
                import math
                safe_dict = {
                    "__builtins__": {}, "sin": math.sin, "cos": math.cos,
                    "exp": math.exp, "log": math.log, "sqrt": math.sqrt,
                    "abs": abs, "pi": math.pi, "e": math.e, "pow": pow
                }
                result = float(eval(obj_expr, safe_dict, {}))
            
            # 如果求最大值，返回负值
            return result if minimize else -result

        result = opt.minimize(objective, initial_point, method='Newton-CG',
                             options={'maxiter': max_iterations, 'gtol': tolerance})

        # 如果求最大值，需要将结果的函数值取负
        optimal_value = float(result.fun) if minimize else -float(result.fun)
        optimization_type = "最小值" if minimize else "最大值"
        
        return {
            "optimal_point": result.x.tolist(),
            "optimal_value": optimal_value,
            "success": bool(result.success),
            "iterations": int(result.nit),
            "message": result.message,
            "method": f"牛顿法（求{optimization_type}）",
            "optimization_type": optimization_type
        }

    except Exception as e:
        return f"错误：牛顿法优化失败 - {str(e)}"

@mcp.tool()
def bfgs_optimization(objective_func: str, initial_point: List[float],
                      tolerance: float = 1e-6, max_iterations: int = 100, minimize: bool = True) -> Union[Dict[str, Any], str]:
    """
    拟牛顿法 BFGS 无约束优化
    
    参数:
    - objective_func: 目标函数表达式
    - initial_point: 初始点
    - tolerance: 收敛容差
    - max_iterations: 最大迭代次数
    - minimize: True为求最小值，False为求最大值
    """
    try:
        def objective(x):
            x = np.asarray(x, dtype=float)
            x_val = float(x[0]) if len(x) > 0 else 0.0
            y_val = float(x[1]) if len(x) > 1 else 0.0
            obj_expr = objective_func.replace('x', str(x_val)).replace('y', str(y_val))
            import math
            safe_dict = {
                "__builtins__": {}, "sin": math.sin, "cos": math.cos,
                "exp": math.exp, "log": math.log, "sqrt": math.sqrt,
                "abs": abs, "pi": math.pi, "e": math.e, "pow": pow
            }
            result = float(eval(obj_expr, safe_dict, {}))
            
            # 如果求最大值，返回负值
            return result if minimize else -result

        result = opt.minimize(objective, initial_point, method='BFGS',
                             options={'maxiter': max_iterations, 'gtol': tolerance})

        # 如果求最大值，需要将结果的函数值取负
        optimal_value = float(result.fun) if minimize else -float(result.fun)
        optimization_type = "最小值" if minimize else "最大值"
        
        return {
            "optimal_point": result.x.tolist(),
            "optimal_value": optimal_value,
            "success": bool(result.success),
            "iterations": int(result.nit),
            "message": result.message,
            "method": f"拟牛顿法 BFGS（求{optimization_type}）",
            "optimization_type": optimization_type
        }

    except Exception as e:
        return f"错误：BFGS优化失败 - {str(e)}"

# ==================== 非线性规划 - 有约束优化模块 ====================

@mcp.tool()
def sqp_optimization(objective_func: str, constraints: List[Dict], initial_point: List[float],
                     tolerance: float = 1e-6, max_iterations: int = 100, minimize: bool = True) -> Union[Dict[str, Any], str]:
    """
    顺序二次规划（SQP）优化器（调用 SLSQP 实现）
    
    参数:
    - objective_func: 目标函数表达式
    - constraints: 约束条件列表
    - initial_point: 初始点
    - tolerance: 收敛容差
    - max_iterations: 最大迭代次数
    - minimize: True为求最小值，False为求最大值
    """
    try:
        def objective(x):
            x = np.asarray(x, dtype=float)
            x_val = float(x[0]) if len(x) > 0 else 0.0
            y_val = float(x[1]) if len(x) > 1 else 0.0
            obj_expr = objective_func.replace('x', str(x_val)).replace('y', str(y_val))
            import math
            safe_dict = {
                "__builtins__": {}, "sin": math.sin, "cos": math.cos,
                "exp": math.exp, "log": math.log, "sqrt": math.sqrt,
                "abs": abs, "pi": math.pi, "e": math.e, "pow": pow
            }
            result = float(eval(obj_expr, safe_dict, {}))
            
            # 如果求最大值，返回负值
            return result if minimize else -result

        scipy_constraints = []
        for con in constraints:
            if con['type'] not in ('eq', 'ineq'):
                continue
            expr = con['fun']
            def constraint_func(x, expr=expr):
                x = np.asarray(x, dtype=float)
                x_val = float(x[0]) if len(x) > 0 else 0.0
                y_val = float(x[1]) if len(x) > 1 else 0.0
                expr_eval = expr.replace('x', str(x_val)).replace('y', str(y_val))
                import math
                safe_dict = {
                    "__builtins__": {}, "sin": math.sin, "cos": math.cos,
                    "exp": math.exp, "log": math.log, "sqrt": math.sqrt,
                    "abs": abs, "pi": math.pi, "e": math.e, "pow": pow
                }
                return float(eval(expr_eval, safe_dict, {}))
            scipy_constraints.append({"type": con['type'], "fun": constraint_func})

        result = opt.minimize(objective, initial_point, method='SLSQP',
                             constraints=scipy_constraints,
                             options={'maxiter': max_iterations, 'ftol': tolerance})

        # 如果求最大值，需要将结果的函数值取负
        optimal_value = float(result.fun) if minimize else -float(result.fun)
        optimization_type = "最小值" if minimize else "最大值"
        
        return {
            "optimal_point": result.x.tolist(),
            "optimal_value": optimal_value,
            "success": bool(result.success),
            "iterations": int(result.nit),
            "message": result.message,
            "method": f"顺序二次规划（SLSQP）（求{optimization_type}）",
            "optimization_type": optimization_type
        }

    except Exception as e:
        return f"错误：SQP优化失败 - {str(e)}"

# ==================== 全局优化模块（元启发式算法） ====================

@mcp.tool()
def differential_evolution_optimization(objective_func: str, bounds: List[Tuple[float, float]],
                                        max_iterations: int = 1000, pop_size: int = 15, minimize: bool = True) -> Union[Dict[str, Any], str]:
    """
    差分进化算法求解全局优化问题
    
    参数:
    - objective_func: 目标函数表达式
    - bounds: 变量边界
    - max_iterations: 最大迭代次数
    - pop_size: 种群大小
    - minimize: True为求最小值，False为求最大值
    """
    try:
        def objective(x):
            x = np.asarray(x, dtype=float)
            x_val = float(x[0]) if len(x) > 0 else 0.0
            y_val = float(x[1]) if len(x) > 1 else 0.0
            obj_expr = objective_func.replace('x', str(x_val)).replace('y', str(y_val))
            import math
            safe_dict = {
                "__builtins__": {}, "sin": math.sin, "cos": math.cos,
                "exp": math.exp, "log": math.log, "sqrt": math.sqrt,
                "abs": abs, "pi": math.pi, "e": math.e, "pow": pow
            }
            result = float(eval(obj_expr, safe_dict, {}))
            
            # 如果求最大值，返回负值
            return result if minimize else -result

        result = opt.differential_evolution(objective, bounds, strategy='best1bin',
                                            maxiter=max_iterations, popsize=pop_size)

        # 如果求最大值，需要将结果的函数值取负
        optimal_value = float(result.fun) if minimize else -float(result.fun)
        optimization_type = "最小值" if minimize else "最大值"
        
        return {
            "optimal_point": result.x.tolist(),
            "optimal_value": optimal_value,
            "success": bool(result.success),
            "iterations": int(result.nit),
            "message": result.message,
            "method": f"差分进化算法（求{optimization_type}）",
            "optimization_type": optimization_type
        }

    except Exception as e:
        return f"错误：差分进化优化失败 - {str(e)}"

# ==================== 线性规划 LP ====================

@mcp.tool()
def linear_programming(c: List[float], A_ub: Optional[List[List[float]]] = None,
                       b_ub: Optional[List[float]] = None,
                       A_eq: Optional[List[List[float]]] = None,
                       b_eq: Optional[List[float]] = None,
                       bounds: Optional[List[List[float]]] = None) -> Union[Dict[str, Any], str]:
    """
    线性规划求解器（调用 linprog）
    bounds参数：每个变量的边界，格式为[[下界, 上界], ...]，使用大数值表示无界
    例如：[[0, 1e10], [-1e10, 10], [0, 100]]
    """
    try:
        # 替换无穷上下界为极值，避免 Pydantic 校验失败
        processed_bounds = []
        if bounds is not None:
            for bound in bounds:
                if bound is None or len(bound) != 2:
                    processed_bounds.append([-1e20, 1e20])  # 默认无界
                else:
                    lower = bound[0] if bound[0] is not None and bound[0] > -1e9 else -1e20
                    upper = bound[1] if bound[1] is not None and bound[1] < 1e9 else 1e20
                    processed_bounds.append([lower, upper])

        result = opt.linprog(c=c, A_ub=A_ub, b_ub=b_ub, A_eq=A_eq, b_eq=b_eq,
                             bounds=processed_bounds if bounds else None,
                             method='highs')

        return {
            "optimal_point": result.x.tolist() if result.success else None,
            "optimal_value": float(result.fun) if result.success else None,
            "success": bool(result.success),
            "message": result.message,
            "method": "线性规划（HiGHS）"
        }
    except Exception as e:
        return f"错误：线性规划求解失败 - {str(e)}"

# ==================== 二次规划 QP ====================

@mcp.tool()
def quadratic_programming(P: List[List[float]], q: List[float],
                          G: Optional[List[List[float]]] = None, h: Optional[List[float]] = None,
                          A: Optional[List[List[float]]] = None, b: Optional[List[float]] = None) -> Union[Dict[str, Any], str]:
    """
    二次规划求解器（使用 cvxopt）
    """
    try:
        from cvxopt import matrix, solvers
        import numpy as np
        solvers.options['show_progress'] = False
        
        # 转换为numpy数组然后转换为cvxopt matrix
        P_mat = matrix(np.array(P, dtype=float))
        q_mat = matrix(np.array(q, dtype=float))
        
        G_mat = None
        h_mat = None
        if G is not None and h is not None:
            G_mat = matrix(np.array(G, dtype=float))
            h_mat = matrix(np.array(h, dtype=float))
        
        A_mat = None
        b_mat = None
        if A is not None and b is not None:
            # cvxopt的A矩阵格式：需要明确指定维度
            A_array = np.array(A, dtype=float)
            if A_array.ndim == 1:
                A_array = A_array.reshape(1, -1)
            # 使用正确的matrix构造方式
            A_mat = matrix(A_array.flatten().tolist(), A_array.shape)
            b_mat = matrix(np.array(b, dtype=float))

        result = solvers.qp(P_mat, q_mat, G_mat, h_mat, A_mat, b_mat)
        
        if result['status'] == 'optimal':
            return {
                "optimal_point": [float(v) for v in result['x']],
                "optimal_value": float(result['primal objective']),
                "status": result['status'],
                "success": True,
                "method": "二次规划（cvxopt）"
            }
        else:
            return {
                "optimal_point": None,
                "optimal_value": None,
                "status": result['status'],
                "success": False,
                "method": "二次规划（cvxopt）"
            }
    except Exception as e:
        return f"错误：二次规划求解失败 - {str(e)}"

# ==================== 整数规划 MILP ====================

@mcp.tool()
def mixed_integer_linear_programming(c: List[float], A: List[List[float]], b: List[float],
                                     integer_indices: List[int],
                                     bounds: Optional[List[Tuple[float, float]]] = None) -> Union[Dict[str, Any], str]:
    """
    混合整数线性规划（使用 pulp）
    """
    try:
        import pulp
        n = len(c)
        model = pulp.LpProblem("MILP", pulp.LpMinimize)
        x_vars = [
            pulp.LpVariable(f"x{i}", lowBound=bounds[i][0] if bounds else None,
                            upBound=bounds[i][1] if bounds else None,
                            cat=pulp.LpInteger if i in integer_indices else pulp.LpContinuous)
            for i in range(n)
        ]
        model += pulp.lpSum([c[i] * x_vars[i] for i in range(n)])
        for i in range(len(A)):
            model += pulp.lpSum([A[i][j] * x_vars[j] for j in range(n)]) <= b[i]

        model.solve()
        return {
            "optimal_point": [pulp.value(var) for var in x_vars],
            "optimal_value": pulp.value(model.objective),
            "status": pulp.LpStatus[model.status],
            "method": "混合整数线性规划（pulp）"
        }
    except Exception as e:
        return f"错误：MILP求解失败 - {str(e)}"


# ==================== 增广拉格朗日法 - 有约束非线性优化模块 ====================

@mcp.tool()
def augmented_lagrangian_optimization(objective_func: str, constraints: List[Dict], initial_point: List[float],
                                     tolerance: float = 1e-6, max_iterations: int = 100) -> Union[Dict[str, Any], str]:
    """
    增广拉格朗日法优化（使用 scipy minimize + penalty 形式实现）
    支持等式和不等式约束
    
    constraints格式示例：
    [
        {"type": "eq", "fun": "x + y - 1"},
        {"type": "ineq", "fun": "x - 0.5"}
    ]
    """
    try:
        import numpy as np
        from scipy.optimize import minimize
        
        def eval_expr(expr, x):
            x = np.asarray(x, dtype=float)
            x_val = float(x[0]) if len(x) > 0 else 0.0
            y_val = float(x[1]) if len(x) > 1 else 0.0
            import math
            safe_dict = {
                "__builtins__": {},
                "sin": math.sin, "cos": math.cos, "exp": math.exp,
                "log": math.log, "sqrt": math.sqrt, "abs": abs,
                "pi": math.pi, "e": math.e, "pow": pow
            }
            expr_eval = expr.replace('x', str(x_val)).replace('y', str(y_val))
            return float(eval(expr_eval, safe_dict, {}))

        # 原始目标函数
        def f(x):
            return eval_expr(objective_func, x)

        # 罚函数参数初始化
        mu = 10.0  # 惩罚因子
        lam = np.zeros(len(constraints))  # 拉格朗日乘子初始化
        
        xk = np.array(initial_point, dtype=float)
        
        for k in range(max_iterations):
            def aug_lag_obj(x):
                penalty = 0.0
                for i, con in enumerate(constraints):
                    c_val = eval_expr(con['fun'], x)
                    if con['type'] == 'eq':
                        penalty += lam[i] * c_val + 0.5 * mu * c_val ** 2
                    elif con['type'] == 'ineq':
                        penalty += lam[i] * min(0, c_val) + 0.5 * mu * min(0, c_val) ** 2
                return f(x) + penalty
            
            res = minimize(aug_lag_obj, xk, method='BFGS', options={'gtol': tolerance})
            xk = res.x
            
            # 更新拉格朗日乘子
            for i, con in enumerate(constraints):
                c_val = eval_expr(con['fun'], xk)
                if con['type'] == 'eq':
                    lam[i] = lam[i] + mu * c_val
                elif con['type'] == 'ineq':
                    lam[i] = lam[i] + mu * min(0, c_val)
            
            # 收敛判断（目标函数变化或约束满足度）
            if res.success:
                cons_violations = [abs(eval_expr(con['fun'], xk)) if con['type']=='eq' else max(0, -eval_expr(con['fun'], xk)) for con in constraints]
                if max(cons_violations) < tolerance:
                    break
            
            mu *= 10  # 惩罚因子加大
        
        return {
            "optimal_point": xk.tolist(),
            "optimal_value": f(xk),
            "iterations": k+1,
            "message": res.message,
            "success": res.success,
            "method": "增广拉格朗日法"
        }
    except Exception as e:
        return f"错误：增广拉格朗日优化失败 - {str(e)}"


# ==================== 内点法（trust-constr） - 有约束非线性优化模块 ====================

@mcp.tool()
def interior_point_optimization(objective_func: str, constraints: List[Dict], initial_point: List[float],
                                tolerance: float = 1e-6, max_iterations: int = 100) -> Union[Dict[str, Any], str]:
    """
    内点法优化（trust-constr）
    支持等式和不等式约束
    
    constraints格式同上
    """
    try:
        import numpy as np
        from scipy.optimize import minimize, NonlinearConstraint
        
        def eval_expr(expr, x):
            x = np.asarray(x, dtype=float)
            x_val = float(x[0]) if len(x) > 0 else 0.0
            y_val = float(x[1]) if len(x) > 1 else 0.0
            import math
            safe_dict = {
                "__builtins__": {},
                "sin": math.sin, "cos": math.cos,
                "exp": math.exp, "log": math.log,
                "sqrt": math.sqrt, "abs": abs,
                "pi": math.pi, "e": math.e, "pow": pow
            }
            expr_eval = expr.replace('x', str(x_val)).replace('y', str(y_val))
            return float(eval(expr_eval, safe_dict, {}))

        def objective(x):
            return eval_expr(objective_func, x)
        
        nl_constraints = []
        for con in constraints:
            if con['type'] == 'eq':
                nl_constraints.append(NonlinearConstraint(lambda x, expr=con['fun']: eval_expr(expr, x), 0, 0))
            elif con['type'] == 'ineq':
                # scipy定义不等式为 fun(x) >= lb; 所以用 lb=0
                nl_constraints.append(NonlinearConstraint(lambda x, expr=con['fun']: eval_expr(expr, x), 0, np.inf))

        result = minimize(objective, initial_point, method='trust-constr',
                          constraints=nl_constraints,
                          options={'gtol': tolerance, 'maxiter': max_iterations})

        return {
            "optimal_point": result.x.tolist(),
            "optimal_value": float(result.fun),
            "success": bool(result.success),
            "iterations": int(result.nit),
            "message": result.message,
            "method": "内点法 (trust-constr)"
        }

    except Exception as e:
        return f"错误：内点法优化失败 - {str(e)}"


# ==================== Zoutendijk可行方向法 - 有约束非线性优化模块 ====================

@mcp.tool()
def zoutendijk_feasible_direction_optimization(objective_func: str, constraints: List[Dict], initial_point: List[float],
                                              tolerance: float = 1e-6, max_iterations: int = 100) -> Union[Dict[str, Any], str]:
    """
    Zoutendijk可行方向法优化
    适用于线性约束的非线性规划问题
    
    constraints格式示例：
    [
        {"type": "eq", "fun": "x + y - 1"},
        {"type": "ineq", "fun": "x - 0.5"}
    ]
    """
    try:
        import numpy as np
        from scipy.optimize import minimize_scalar
        
        def eval_expr(expr, x):
            x = np.asarray(x, dtype=float)
            x_val = float(x[0]) if len(x) > 0 else 0.0
            y_val = float(x[1]) if len(x) > 1 else 0.0
            import math
            safe_dict = {
                "__builtins__": {},
                "sin": math.sin, "cos": math.cos, "exp": math.exp,
                "log": math.log, "sqrt": math.sqrt, "abs": abs,
                "pi": math.pi, "e": math.e, "pow": pow
            }
            expr_eval = expr.replace('x', str(x_val)).replace('y', str(y_val))
            return float(eval(expr_eval, safe_dict, {}))
        
        def gradient_numerical(func, x, h=1e-8):
            """数值梯度计算"""
            grad = np.zeros_like(x)
            for i in range(len(x)):
                x_plus = x.copy()
                x_minus = x.copy()
                x_plus[i] += h
                x_minus[i] -= h
                grad[i] = (func(x_plus) - func(x_minus)) / (2 * h)
            return grad
        
        def objective(x):
            return eval_expr(objective_func, x)
        
        def constraint_value(x, constraint):
            return eval_expr(constraint['fun'], x)
        
        def is_feasible(x, constraints, tol=1e-10):
            """检查点是否可行"""
            for con in constraints:
                val = constraint_value(x, con)
                if con['type'] == 'eq' and abs(val) > tol:
                    return False
                elif con['type'] == 'ineq' and val < -tol:
                    return False
            return True
        
        def find_feasible_direction(x, constraints, grad_f):
            """寻找可行方向"""
            n = len(x)
            
            # 识别活跃约束
            active_constraints = []
            for i, con in enumerate(constraints):
                val = constraint_value(x, con)
                if con['type'] == 'eq' or (con['type'] == 'ineq' and abs(val) < 1e-8):
                    active_constraints.append(i)
            
            if not active_constraints:
                # 无活跃约束，使用负梯度方向
                return -grad_f / np.linalg.norm(grad_f)
            
            # 构建活跃约束的梯度矩阵
            A = []
            for idx in active_constraints:
                con = constraints[idx]
                grad_con = gradient_numerical(lambda y: constraint_value(y, con), x)
                A.append(grad_con)
            
            A = np.array(A)
            
            # 求解线性系统找到可行方向
            # 使用投影方法：d = -grad_f + A^T * lambda
            # 其中 A * d = 0 (对于活跃约束)
            
            try:
                # 使用QR分解求解零空间
                if A.shape[0] > 0:
                    Q, R = np.linalg.qr(A.T)
                    # 零空间基
                    null_space = Q[:, A.shape[0]:]
                    if null_space.shape[1] > 0:
                        # 在零空间中投影负梯度
                        proj_grad = null_space @ (null_space.T @ (-grad_f))
                        if np.linalg.norm(proj_grad) > 1e-12:
                            return proj_grad / np.linalg.norm(proj_grad)
                
                # 如果无法找到下降方向，尝试使用最小二乘解
                try:
                    lambda_opt = np.linalg.lstsq(A, grad_f, rcond=None)[0]
                    d = -grad_f + A.T @ lambda_opt
                    if np.linalg.norm(d) > 1e-12:
                        return d / np.linalg.norm(d)
                except:
                    pass
                    
            except np.linalg.LinAlgError:
                pass
            
            # 如果所有方法都失败，返回零向量（收敛）
            return np.zeros(n)
        
        def line_search(x, direction, constraints, alpha_max=1.0):
            """线搜索确定步长"""
            if np.linalg.norm(direction) < 1e-12:
                return 0.0
            
            # 确定最大可行步长
            alpha_feasible = alpha_max
            for con in constraints:
                if con['type'] == 'ineq':
                    # 对于不等式约束 g(x) >= 0
                    # 需要 g(x + alpha*d) >= 0
                    x_test = x + 1e-6 * direction
                    grad_con = gradient_numerical(lambda y: constraint_value(y, con), x)
                    denom = np.dot(grad_con, direction)
                    if denom < -1e-12:  # 方向会违反约束
                        val = constraint_value(x, con)
                        alpha_bound = val / (-denom)
                        alpha_feasible = min(alpha_feasible, alpha_bound * 0.99)
            
            # Armijo线搜索
            c1 = 1e-4
            alpha = min(alpha_feasible, 1.0)
            grad_f = gradient_numerical(objective, x)
            descent_condition = np.dot(grad_f, direction)
            
            if descent_condition >= 0:  # 不是下降方向
                return 0.0
            
            f_x = objective(x)
            
            for _ in range(20):  # 最多20次回退
                if alpha < 1e-12:
                    break
                x_new = x + alpha * direction
                if is_feasible(x_new, constraints):
                    f_new = objective(x_new)
                    if f_new <= f_x + c1 * alpha * descent_condition:
                        return alpha
                alpha *= 0.5
            
            return 0.0
        
        # 主算法循环
        x = np.array(initial_point, dtype=float)
        history = [x.copy()]
        
        if not is_feasible(x, constraints):
            return {
                "optimal_point": None,
                "optimal_value": None,
                "success": False,
                "iterations": 0,
                "message": "初始点不可行",
                "method": "Zoutendijk可行方向法"
            }
        
        for iteration in range(max_iterations):
            # 计算目标函数梯度
            grad_f = gradient_numerical(objective, x)
            
            # 检查一阶最优性条件
            if np.linalg.norm(grad_f) < tolerance:
                break
            
            # 寻找可行方向
            direction = find_feasible_direction(x, constraints, grad_f)
            
            # 检查是否找到有效方向
            if np.linalg.norm(direction) < tolerance:
                break
            
            # 线搜索确定步长
            alpha = line_search(x, direction, constraints)
            
            if alpha < tolerance:
                break
            
            # 更新点
            x_new = x + alpha * direction
            
            # 检查收敛性
            if np.linalg.norm(x_new - x) < tolerance:
                x = x_new
                break
            
            x = x_new
            history.append(x.copy())
        
        return {
            "optimal_point": x.tolist(),
            "optimal_value": float(objective(x)),
            "success": True,
            "iterations": iteration + 1,
            "message": "Zoutendijk可行方向法收敛",
            "method": "Zoutendijk可行方向法",
            "history": [point.tolist() for point in history]
        }
        
    except Exception as e:
        return f"错误：Zoutendijk可行方向法优化失败 - {str(e)}"


# ==================== 智能推荐器模块 ====================

@mcp.tool()
def recommend_optimizer(
    problem_type: str,
    has_constraints: bool,
    has_integer_vars: bool,
    is_blackbox: bool,
    dimension: int,
    is_convex: Optional[bool] = None
) -> str:
    """
    基于问题特征推荐最适合的优化算法
    """
    if has_integer_vars:
        if problem_type == "linear":
            return "推荐算法：混合整数线性规划 MILP（pulp）"
        else:
            return "推荐算法：遗传算法 GA 或 模拟退火 SA（适用于整数/非线性问题）"

    if is_blackbox:
        if dimension <= 100:
            return "推荐算法：差分进化 DE（适用于黑盒/非凸问题）"
        else:
            return "推荐算法：粒子群优化 PSO 或 遗传算法 GA（高维黑盒优化）"

    if problem_type == "linear":
        return "推荐算法：线性规划 linprog（HiGHS）"

    if problem_type == "quadratic":
        return "推荐算法：二次规划 QP（cvxopt）"

    if has_constraints:
        return "推荐算法：SQP 或 增广拉格朗日法 或 内点法 或 Zoutendijk可行方向法（支持等式/不等式约束）"

    if dimension > 100:
        return "推荐算法：L-BFGS 或 PSO（适用于高维优化）"

    if not has_constraints and problem_type == "nonlinear":
        return "推荐算法：BFGS 或 牛顿法（快速局部收敛）"

    # 新增的推荐算法
    if is_convex is not None:
        if is_convex:
            return "推荐算法：信赖域方法（适用于凸优化问题）"
        else:
            return "推荐算法：梯度投影法 或 可行方向法（适用于非凸优化问题）"

    return "推荐算法：默认使用 BFGS 或 SQP；若不确定建议使用 DE 或 GA"



# ==================== 可视化模块 ====================

@mcp.tool()
def visualize_optimization_problem(
    objective_func: str,
    constraints: Optional[List[Dict[str, str]]] = None,
    optimal_point: Optional[List[float]] = None,
    x_range: Optional[List[float]] = None,
    y_range: Optional[List[float]] = None,
    save_path: Optional[str] = None
) -> Union[Dict[str, Any], str]:
    """
    可视化约束优化问题：目标函数等高线 + 约束条件 + 最优解
    """
    try:
        import matplotlib.pyplot as plt
        import numpy as np
        
        # 将Python表达式转换为LaTeX格式
        def python_to_latex(expr):
            """将Python数学表达式转换为LaTeX格式"""
            latex_expr = expr
            import re
            # 替换不等号
            latex_expr = re.sub(r'<=', r'\\leq', latex_expr)
            latex_expr = re.sub(r'>=', r'\\geq', latex_expr)
            latex_expr = re.sub(r'!=', r'\\neq', latex_expr)
            # 替换幂运算
            latex_expr = re.sub(r'\*\*([0-9]+)', r'^{\1}', latex_expr)
            latex_expr = re.sub(r'\*\*\(([^)]+)\)', r'^{\1}', latex_expr)
            # 替换乘法符号
            latex_expr = re.sub(r'\*', r'\\cdot ', latex_expr)
            # 处理括号中的表达式
            latex_expr = re.sub(r'\(([^)]+)\)', r'(\1)', latex_expr)
            return latex_expr
        
        # 设置字体
        import matplotlib.font_manager as fm
        try:
            # 尝试使用系统中文字体
            chinese_fonts = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans', 'Arial Unicode MS']
            available_fonts = [f.name for f in fm.fontManager.ttflist]
            
            for font in chinese_fonts:
                if font in available_fonts:
                    plt.rcParams['font.sans-serif'] = [font]
                    break
            else:
                plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
            
            plt.rcParams['axes.unicode_minus'] = False
            plt.rcParams['font.family'] = plt.rcParams['font.sans-serif']
            plt.rcParams['mathtext.fontset'] = 'stix'
            plt.rcParams['mathtext.rm'] = 'Times New Roman'
            plt.rcParams['mathtext.it'] = 'Times New Roman:italic'
            plt.rcParams['mathtext.bf'] = 'Times New Roman:bold'
            plt.rcParams['font.size'] = 12
        except Exception:
            plt.rcParams['font.family'] = 'DejaVu Sans'
        
        # 设置范围
        x_min, x_max = x_range if x_range else [-5, 5]
        y_min, y_max = y_range if y_range else [-5, 5]
        
        # 创建网格
        x = np.linspace(x_min, x_max, 400)
        y = np.linspace(y_min, y_max, 400)
        X, Y = np.meshgrid(x, y)
        
        # 计算目标函数值
        def eval_objective(x_val, y_val):
            import math
            safe_dict = {
                "__builtins__": {},
                "sin": math.sin, "cos": math.cos, "exp": math.exp,
                "log": math.log, "sqrt": math.sqrt, "abs": abs,
                "pi": math.pi, "e": math.e, "pow": pow
            }
            expr = objective_func.replace('x', str(x_val)).replace('y', str(y_val))
            return eval(expr, safe_dict, {})
        
        Z = np.zeros_like(X)
        for i in range(X.shape[0]):
            for j in range(X.shape[1]):
                try:
                    Z[i, j] = eval_objective(X[i, j], Y[i, j])
                except:
                    Z[i, j] = np.nan
        
        # 创建图形
        fig, ax = plt.subplots(1, 1, figsize=(12, 10))
        
        # 绘制目标函数等高线
        contour = ax.contour(X, Y, Z, levels=20, colors='lightblue', alpha=0.6)
        ax.clabel(contour, inline=True, fontsize=8)
        
        # 绘制约束条件
        if constraints:
            colors = ['red', 'orange', 'green', 'purple', 'brown']
            for i, constraint in enumerate(constraints):
                color = colors[i % len(colors)]
                constraint_type = constraint.get('type', 'eq')
                constraint_func = constraint.get('fun', '')
                
                if constraint_type == 'eq':
                    # 等式约束：绘制零等高线
                    def eval_constraint(x_val, y_val):
                        import math
                        safe_dict = {
                            "__builtins__": {},
                            "sin": math.sin, "cos": math.cos, "exp": math.exp,
                            "log": math.log, "sqrt": math.sqrt, "abs": abs,
                            "pi": math.pi, "e": math.e, "pow": pow
                        }
                        expr = constraint_func.replace('x', str(x_val)).replace('y', str(y_val))
                        return eval(expr, safe_dict, {})
                    
                    C = np.zeros_like(X)
                    for ii in range(X.shape[0]):
                        for jj in range(X.shape[1]):
                            try:
                                C[ii, jj] = eval_constraint(X[ii, jj], Y[ii, jj])
                            except:
                                C[ii, jj] = np.nan
                    
                    ax.contour(X, Y, C, levels=[0], colors=[color], linewidths=2)
                    # 使用display字段或转换约束函数为LaTeX格式
                    display_text = constraint.get('display', constraint_func)
                    latex_constraint = python_to_latex(display_text)
                    ax.plot([], [], color=color, linewidth=2, label=f'等式约束 {i+1}: ${latex_constraint}$')
                    
                elif constraint_type == 'ineq':
                    # 不等式约束：绘制可行域
                    def eval_constraint(x_val, y_val):
                        import math
                        safe_dict = {
                            "__builtins__": {},
                            "sin": math.sin, "cos": math.cos, "exp": math.exp,
                            "log": math.log, "sqrt": math.sqrt, "abs": abs,
                            "pi": math.pi, "e": math.e, "pow": pow
                        }
                        expr = constraint_func.replace('x', str(x_val)).replace('y', str(y_val))
                        return eval(expr, safe_dict, {})
                    
                    C = np.zeros_like(X)
                    for ii in range(X.shape[0]):
                        for jj in range(X.shape[1]):
                            try:
                                C[ii, jj] = eval_constraint(X[ii, jj], Y[ii, jj])
                            except:
                                C[ii, jj] = np.nan
                    
                    # 绘制边界线
                    ax.contour(X, Y, C, levels=[0], colors=[color], linewidths=2, linestyles='--')
                    # 填充可行域
                    ax.contourf(X, Y, C, levels=[0, np.inf], colors=[color], alpha=0.2)
                    # 使用display字段或转换约束函数为LaTeX格式
                    display_text = constraint.get('display', f'{constraint_func} \\geq 0')
                    latex_constraint = python_to_latex(display_text)
                    ax.plot([], [], color=color, linewidth=2, linestyle='--', label=f'不等式约束 {i+1}: ${latex_constraint}$')
        
        # 标记最优解
        if optimal_point and len(optimal_point) >= 2:
            ax.plot(optimal_point[0], optimal_point[1], 'ro', markersize=12, label=f'最优解: $({optimal_point[0]:.3f}, {optimal_point[1]:.3f})$')
            
            # 计算最优值
            try:
                optimal_value = eval_objective(optimal_point[0], optimal_point[1])
                ax.text(optimal_point[0], optimal_point[1] + 0.2, f'最优值: ${optimal_value:.6f}$', 
                        ha='center', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.8))
            except:
                pass
        
        # 设置图形属性
        ax.set_xlim(x_min, x_max)
        ax.set_ylim(y_min, y_max)
        ax.set_xlabel('$x$', fontsize=14)
        ax.set_ylabel('$y$', fontsize=14)
        latex_func = python_to_latex(objective_func)
        ax.set_title(f'约束优化问题可视化\n目标函数: ${latex_func}$', fontsize=16)
        ax.grid(True, alpha=0.3)
        ax.legend(fontsize=10)
        
        # 保存图形
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        plt.tight_layout()
        plt.show()
        
        return {
            "status": "success",
            "message": "可视化完成",
            "save_path": save_path if save_path else "未保存"
        }
        
    except Exception as e:
        return f"错误：可视化失败 - {str(e)}"

@mcp.tool()
def visualize_convergence_history(
    history: List[List[float]],
    objective_values: Optional[List[float]] = None,
    save_path: Optional[str] = None
) -> Union[Dict[str, Any], str]:
    """
    可视化优化收敛过程：迭代轨迹 + 目标函数值变化
    返回base64编码的图片
    """
    try:
        import matplotlib.pyplot as plt
        import numpy as np
        import base64
        import io
        
        # 设置字体
        import matplotlib.font_manager as fm
        try:
            # 尝试使用系统中文字体
            chinese_fonts = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans', 'Arial Unicode MS']
            available_fonts = [f.name for f in fm.fontManager.ttflist]
            
            for font in chinese_fonts:
                if font in available_fonts:
                    plt.rcParams['font.sans-serif'] = [font]
                    break
            else:
                plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
            
            plt.rcParams['axes.unicode_minus'] = False
            plt.rcParams['font.family'] = plt.rcParams['font.sans-serif']
            plt.rcParams['mathtext.fontset'] = 'stix'
            plt.rcParams['mathtext.rm'] = 'Times New Roman'
            plt.rcParams['mathtext.it'] = 'Times New Roman:italic'
            plt.rcParams['mathtext.bf'] = 'Times New Roman:bold'
            plt.rcParams['font.size'] = 12
        except Exception:
            plt.rcParams['font.family'] = 'DejaVu Sans'
        
        history = np.array(history)
        
        if history.shape[1] == 2:
            # 二维优化轨迹
            fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 6))
            
            # 左图：迭代轨迹
            ax1.plot(history[:, 0], history[:, 1], 'b-o', markersize=4, linewidth=1.5, label='优化轨迹')
            ax1.plot(history[0, 0], history[0, 1], 'go', markersize=8, label='起始点')
            ax1.plot(history[-1, 0], history[-1, 1], 'ro', markersize=8, label='最终点')
            ax1.set_xlabel('$x$', fontsize=14)
            ax1.set_ylabel('$y$', fontsize=14)
            ax1.set_title('优化轨迹', fontsize=16)
            ax1.grid(True, alpha=0.3)
            ax1.legend()
            
            # 右图：目标函数值变化
            if objective_values:
                ax2.plot(range(len(objective_values)), objective_values, 'r-o', markersize=4, linewidth=1.5)
                ax2.set_xlabel('迭代次数', fontsize=14)
                ax2.set_ylabel('目标函数值', fontsize=14)
                ax2.set_title('收敛过程', fontsize=16)
                ax2.grid(True, alpha=0.3)
                ax2.set_yscale('log')
            else:
                ax2.text(0.5, 0.5, '无目标函数数据', ha='center', va='center', transform=ax2.transAxes)
                ax2.set_title('目标函数值变化', fontsize=16)
        
        else:
            # 高维优化：只显示目标函数值变化
            fig, ax = plt.subplots(1, 1, figsize=(10, 6))
            
            if objective_values:
                ax.plot(range(len(objective_values)), objective_values, 'r-o', markersize=4, linewidth=1.5)
                ax.set_xlabel('迭代次数', fontsize=14)
                ax.set_ylabel('目标函数值', fontsize=14)
                ax.set_title('高维优化收敛过程', fontsize=16)
                ax.grid(True, alpha=0.3)
                ax.set_yscale('log')
            else:
                ax.text(0.5, 0.5, '无目标函数数据', ha='center', va='center', transform=ax.transAxes)
                ax.set_title('收敛过程', fontsize=16)
        
        plt.tight_layout()
        
        # 保存到内存并转换为base64
        buffer = io.BytesIO()
        plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
        buffer.seek(0)
        image_base64 = base64.b64encode(buffer.getvalue()).decode('utf-8')
        buffer.close()
        
        # 可选：保存到文件
        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
        
        plt.close(fig)  # 关闭图形以释放内存
        
        return {
            "status": "success",
            "message": "收敛过程可视化完成",
            "iterations": len(history),
            "image_base64": image_base64,
            "save_path": save_path if save_path else "未保存"
        }
        
    except Exception as e:
        return f"错误：收敛过程可视化失败 - {str(e)}"



@mcp.tool()
def solve_optimization_problem(
    objective_func: str,
    initial_point: Optional[List[float]] = None,
    constraints: Optional[List[Dict[str, str]]] = None,
    bounds: Optional[List[Tuple[float, float]]] = None,
    problem_type: str = "nonlinear",  # "linear", "quadratic", "nonlinear"
    has_integer_vars: bool = False,
    is_blackbox: bool = False
) -> Union[Dict[str, Any], str]:
    """
    智能求解优化问题：根据输入自动推荐算法并求解
    """
    try:
        has_constraints = bool(constraints)
        dimension = len(initial_point) if initial_point else 2
        recommended = recommend_optimizer(
            problem_type=problem_type,
            has_constraints=has_constraints,
            has_integer_vars=has_integer_vars,
            is_blackbox=is_blackbox,
            dimension=dimension,
            is_convex=True  # 可选
        )

        # 统一调度对应求解函数
        if has_integer_vars and problem_type == "linear":
            # MILP
            c = [float(term) for term in objective_func.strip().split()]  # 例如：objective_func="1 2 3"
            A = [[1]*len(c)]  # dummy
            b = [1000]
            return {
                "推荐算法": recommended,
                "求解结果": mixed_integer_linear_programming(c=c, A=A, b=b,
                                                           integer_indices=list(range(len(c))),
                                                           bounds=bounds)
            }

        elif problem_type == "linear":
            # Linear Programming
            c = [float(term) for term in objective_func.strip().split()]
            A_ub = [[1]*len(c)]
            b_ub = [1000]
            return {
                "推荐算法": recommended,
                "求解结果": linear_programming(c=c, A_ub=A_ub, b_ub=b_ub, bounds=bounds)
            }

        elif has_constraints:
            if is_blackbox:
                return {
                    "推荐算法": recommended,
                    "求解结果": differential_evolution_optimization(objective_func=objective_func, bounds=bounds, minimize=True)
                }
            else:
                return {
                    "推荐算法": recommended,
                    "求解结果": sqp_optimization(objective_func=objective_func, constraints=constraints, initial_point=initial_point)
                }

        elif not has_constraints:
            return {
                "推荐算法": recommended,
                "求解结果": bfgs_optimization(objective_func=objective_func, initial_point=initial_point)
            }

        else:
            return {
                "推荐算法": "默认算法：差分进化",
                "求解结果": differential_evolution_optimization(objective_func=objective_func, bounds=bounds, minimize=True)
            }

    except Exception as e:
        return f"错误：自动求解失败 - {str(e)}"


if __name__ == "__main__":
    # 启动MCP服务器，使用stdio传输协议
    mcp.run(transport="stdio")
